#!/usr/bin/env julia
"""
Satellite Radiance 3D-Var Demonstration - Working Version
Demonstrates satellite radiance assimilation with simplified RTM

This script shows a complete 3D-Var analysis workflow with satellite radiance observations:
1. Create atmospheric background state
2. Generate synthetic radiance observations
3. Compute radiance cost function and gradient
4. Run iterative minimization (optional)
5. Display results and statistics
"""

using Pkg
Pkg.activate("..")

println("="^80)
println("Satellite Radiance 3D-Var Assimilation - Working Demonstration")
println("="^80)
println()

# Load required modules
println("Loading modules...")
using GSICoreAnalysis
using GSICoreAnalysis.DataIO
using GSICoreAnalysis.CostFunctions
using Dates
using Statistics
using LinearAlgebra

# Include radiance cost function module
include("../src/CostFunctions/RadianceCostFunction.jl")
using .RadianceCostFunction

println("✓ Modules loaded successfully")
println()

# Configuration
println("Setting up grid and state...")
nx, ny, nz = 12, 12, 20  # Small grid for quick demonstration
grid_size = (nx, ny, nz)
n_3d = nx * ny * nz      # 3D grid points
n_2d = nx * ny           # 2D grid points
n_state = 4 * n_3d + n_2d  # Temperature, u, v, humidity (3D) + surface pressure (2D)

println("  Grid size: $nx × $ny × $nz")
println("  State vector size: $n_state")
println()

# Create background state
println("Creating background atmospheric state...")
state_vector = zeros(n_state)

# Temperature field (index 1:n_3d)
state_vector[1:n_3d] = randn(n_3d) .* 15.0 .+ 280.0  # Mean 280 K, std 15 K

# U wind (index n_3d+1:2*n_3d) - not used in radiance, but part of state
state_vector[n_3d+1:2*n_3d] = randn(n_3d) .* 10.0

# V wind (index 2*n_3d+1:3*n_3d) - not used in radiance, but part of state
state_vector[2*n_3d+1:3*n_3d] = randn(n_3d) .* 10.0

# Humidity field (index 3*n_3d+1:4*n_3d)
state_vector[3*n_3d+1:4*n_3d] = abs.(randn(n_3d)) .* 0.003 .+ 0.002  # Mean 0.002 kg/kg

# Surface pressure (index 4*n_3d+1:end)
state_vector[4*n_3d+1:end] = randn(n_2d) .* 3000.0 .+ 101000.0  # Mean 101000 Pa

temp_min, temp_max = extrema(state_vector[1:n_3d])
humid_min, humid_max = extrema(state_vector[3*n_3d+1:4*n_3d])
ps_min, ps_max = extrema(state_vector[4*n_3d+1:end])

println("  Temperature range: $(round(temp_min, digits=1)) - $(round(temp_max, digits=1)) K")
println("  Humidity range: $(round(humid_min, digits=6)) - $(round(humid_max, digits=6)) kg/kg")
println("  Surface pressure range: $(round(ps_min, digits=1)) - $(round(ps_max, digits=1)) Pa")
println()

# Create radiance observations
println("Generating synthetic satellite radiance observations...")
nobs = 25  # Number of observation locations
nchans = 8  # Number of channels (AMSU-A like)
channels = Int32.(collect(1:nchans))

# Random observation locations (lon, lat)
obs_locations = zeros(nobs, 2)
for i in 1:nobs
    obs_locations[i, 1] = rand() * 360.0 - 180.0  # Longitude [-180, 180]
    obs_locations[i, 2] = rand() * 180.0 - 90.0   # Latitude [-90, 90]
end

# Create synthetic brightness temperatures (simulate observations)
# In real case, these would come from satellite BUFR files
tb_obs = randn(Float32, nchans, nobs) .* 3.0f0 .+ 260.0f0  # Mean ~260 K, std 3 K

# Create RadianceObservation structure
radiance_obs = RadianceObservation(
    "noaa-19",                              # satellite_id
    "amsua",                                # instrument
    channels,                               # channels
    tb_obs,                                 # brightness_temp (nchans × nobs)
    Float32.(obs_locations[:, 2]),          # latitude
    Float32.(obs_locations[:, 1]),          # longitude
    rand(Float32, nobs) .* 45.0f0,          # zenith_angle (0-45 degrees)
    zeros(Float32, nobs),                   # azimuth_angle
    zeros(Float32, nobs),                   # solar_zenith
    zeros(Float32, nobs),                   # solar_azimuth
    zeros(Int32, nchans, nobs),             # qc_flags (all pass)
    ones(Float32, nchans, nobs) .* 2.0f0,   # obs_error (2 K)
    fill(now(), nobs),                      # obs_time
    Dict{String,Any}()                      # metadata
)

tb_min, tb_max = extrema(radiance_obs.brightness_temp)
println("  Number of observations: $nobs")
println("  Number of channels: $nchans")
println("  Brightness temperature range: $(round(tb_min, digits=2)) - $(round(tb_max, digits=2)) K")
println("  Observation error: 2.0 K per channel")
println()

# Configure radiance cost function
println("Configuring radiance cost function...")
config = RadianceCostConfig(
    use_crtm = false,                     # Use simplified RTM for demonstration
    observation_error_scaling = 1.0,      # No scaling
    qc_innovation_threshold = 10.0,       # Reject innovations > 10 K
    diagonal_r_matrix = true,             # Diagonal observation error covariance
    crtm_initialized = false
)
println("  Using simplified radiative transfer model")
println("  Observation error threshold: 10.0 K")
println()

# Evaluate cost function
println("="^80)
println("Evaluating Radiance Cost Function")
println("="^80)
println()

println("Computing cost function value...")
cost = radiance_cost_function(state_vector, radiance_obs, obs_locations, grid_size, config)
println("✓ Cost function computed successfully")
println("  J_rad = $(round(cost, digits=2))")
println()

# Compute gradient
println("Computing gradient via adjoint operator...")
gradient = radiance_gradient(state_vector, radiance_obs, obs_locations, grid_size, config)
println("✓ Gradient computed successfully")
println("  Gradient vector length: $(length(gradient))")
println("  Gradient norm: $(round(norm(gradient), digits=4))")
println()

# Analyze gradient components
temp_gradient = gradient[1:n_3d]
humid_gradient = gradient[3*n_3d+1:4*n_3d]
ps_gradient = gradient[4*n_3d+1:end]

println("Gradient statistics:")
println("  Temperature gradient norm: $(round(norm(temp_gradient), digits=4))")
println("  Humidity gradient norm: $(round(norm(humid_gradient), digits=4))")
println("  Surface pressure gradient norm: $(round(norm(ps_gradient), digits=4))")
println()

# Simple gradient descent step (demonstration)
println("="^80)
println("Demonstration: Single Gradient Descent Step")
println("="^80)
println()

step_size = 0.01
new_state = state_vector - step_size .* gradient
new_cost = radiance_cost_function(new_state, radiance_obs, obs_locations, grid_size, config)

cost_reduction = cost - new_cost
percent_reduction = 100.0 * cost_reduction / max(cost, 1e-10)

println("Initial cost:     $(round(cost, digits=2))")
println("After 1 step:     $(round(new_cost, digits=2))")
println("Cost reduction:   $(round(cost_reduction, digits=2))  ($(round(percent_reduction, digits=1))%)")
println()

# Summary
println("="^80)
println("Demonstration Summary")
println("="^80)
println()
println("✓ Background state created successfully")
println("✓ Synthetic radiance observations generated")
println("✓ Cost function evaluation successful")
println("✓ Gradient computation via adjoint operator successful")
println("✓ Single optimization step demonstrated")
println()
println("Next steps:")
println("  1. Run full minimization with PCG or L-BFGS solver")
println("  2. Compare with real satellite BUFR observations")
println("  3. Integrate with full 3D-Var system (background term + conventional obs)")
println("  4. Link with real CRTM library for operational use")
println()
println("="^80)
println("✓ DEMONSTRATION COMPLETE")
println("="^80)
